%% US Public Debt and Safe Asset Market Power
%% Jason Choi, Rishabh Kirpalani, and Diego Perez
%% Nov 24, 2024

%% Solve N=2 Cournot Equilibrium 

%----------------------------------------------------------------
% 0. Housekeeping
%----------------------------------------------------------------

close all

%----------------------------------------------------------------
% 1. Defining variables
%----------------------------------------------------------------

// Endogenous Variables
var bus brw rb rkstar rk krw_star kus_star krw kus kstar k wstar w c_rw c_us drdb vrw vus y dMrwdb;

// Exogenous Variables
var nnu oomega A Astar;

// Shocks
varexo eps_nnu eps_oomega eps_A;

// Paramters
parameters ggamma bbeta eeta llambda aalpha Astarbar Abar iiota iiota_star ddelta_rw ddelta_us N_2
nnu_bar oomega_bar rrho_nnu rrho_oomega ssigma_nnu ssigma_oomega rrho_A ssigma_A rrho_Astar ssigma_Astar
brw_me2 rb_me2 rkstar_me2 rk_me2 krw_star_me2 kus_star_me2 krw_me2 kus_me2 kstar_me2 k_me2
capKstar_me2 capK_me2 wstar_me2 w_me2 crw_me2 cus_me2 vrw_me2 vus_me2 dMrwdb_me2 y_me2;

%----------------------------------------------------------------
% 2. Calibration
%----------------------------------------------------------------

% // Parameters
ggamma = 2;
bbeta = 0.9886;
eeta = 0.545;
llambda = 1;
aalpha = 0.3;
Astarbar = 0.9254;
Abar = 0.8154;
iiota = 0.9070;
iiota_star = 0.7939;
ddelta_rw = 0.1;
ddelta_us = 0.1;
N_2 = 2;
nnu_bar = 0.0042;
oomega_bar = 0.0063;
rrho_nnu = 0.99;
ssigma_nnu = 0.01;
rrho_oomega = 0.95;
ssigma_oomega = 0.3;
rrho_A = 0.95;
ssigma_A = 0.02;
rrho_Astar = rrho_A;
ssigma_Astar = ssigma_A;

% // Analytic Steady State (Monopoly Equilibrium, N=2)
brw_me2 = (oomega_bar/(N_2^llambda*nnu_bar*(1+1/N_2*(eeta-1))))^(1/(eeta-1-llambda));
bus_me2 = 1/N_2*(oomega_bar/(N_2^llambda*nnu_bar*(1+1/N_2*(eeta-1))))^(1/(eeta-1-llambda));
rb_me2 = 1/bbeta - nnu_bar*(brw_me2)^(eeta-1) - 1;
rkstar_me2 = 1/bbeta + ddelta_rw - 1;
rk_me2 = 1/bbeta + ddelta_us - 1;
krw_star_me2 = ((aalpha*(1-iiota_star)*Astarbar*((1/bbeta+ddelta_rw-1)/(aalpha*iiota_star*Astarbar))^((aalpha*(1-iiota_star)-1)/(aalpha*(1-iiota_star))))/(1/bbeta+ddelta_us-1))^((aalpha*(1-iiota_star))/(1-aalpha));
kus_star_me2 = ((1/bbeta+ddelta_rw-1)/(aalpha*iiota_star*Astarbar))^(1/(aalpha*(1-iiota_star)))*krw_star_me2^((1-iiota_star*aalpha)/(aalpha*(1-iiota_star)));
krw_me2 = ((aalpha*iiota*Abar*((1/bbeta+ddelta_us-1)/(aalpha*(1-iiota)*Abar))^((aalpha*iiota-1)/(aalpha*iiota)))/(1/bbeta+ddelta_rw-1))^((aalpha*iiota)/(1-aalpha));
kus_me2 = ((1/bbeta+ddelta_us-1)/(aalpha*(1-iiota)*Abar))^(1/(aalpha*iiota))*krw_me2^((1-(1-iiota)*aalpha)/(aalpha*iiota));
kstar_me2 = krw_star_me2 + krw_me2;
k_me2 = kus_star_me2 + kus_me2;
capKstar_me2 = krw_star_me2^iiota_star*kus_star_me2^(1-iiota_star);
capK_me2 = krw_me2^(1-iiota)*kus_me2^iiota;
wstar_me2 = Astarbar*(1-aalpha)*(capKstar_me2)^aalpha;
w_me2 = Abar*(1-aalpha)*(capK_me2)^aalpha;
crw_me2 = wstar_me2 + (rkstar_me2-ddelta_rw)*kstar_me2 + nnu_bar/eeta*(brw_me2)^eeta + rb_me2*brw_me2;
cus_me2 = w_me2 + (rk_me2-ddelta_us)*k_me2 - oomega_bar/(1+llambda)*(bus_me2)^(1+llambda) - rb_me2*(bus_me2);
drdb_me2 = -nnu_bar*(eeta-1)*brw_me2^(eeta-2);
dMrwdb_me2 = 0;
nnu_me2 = nnu_bar;
oomega_me2 = oomega_bar;
A_me2 = Abar;
Astar_me2 = Astarbar;
vrw_me2 = crw_me2^(1-ggamma)/(1-ggamma)/(1-bbeta);
vus_me2 = cus_me2^(1-ggamma)/(1-ggamma)/(1-bbeta);
y_me2 = A_me2*(kus_me2^iiota*krw_me2^(1-iiota))^aalpha;

%----------------------------------------------------------------
% 3. Model
%----------------------------------------------------------------

model;

brw = bus*N_2;

c_rw^(-ggamma) = bbeta*c_rw(+1)^(-ggamma)*(nnu*brw^(eeta-1)+1+rb);
c_rw^(-ggamma) = bbeta*c_rw(+1)^(-ggamma)*(1-ddelta_rw+rkstar);
c_rw + kstar + brw = wstar + (1-ddelta_rw+rkstar(-1))*kstar(-1) + nnu(-1)/eeta*(brw(-1))^eeta + (1+rb(-1))*brw(-1);

c_us^(-ggamma) = bbeta*c_us(+1)^(-ggamma)*(oomega*(bus)^(llambda)+(1+rb)+drdb*bus);
c_us^(-ggamma) = bbeta*c_us(+1)^(-ggamma)*(1-ddelta_us+rk);
c_us + k - bus = w + (1-ddelta_us+rk(-1))*k(-1) - oomega(-1)/(1+llambda)*(bus(-1))^(1+llambda) - (1+rb(-1))*bus(-1);

rk = Astar*aalpha*(1-iiota_star)*krw_star^(aalpha*iiota_star)*kus_star^(aalpha*(1-iiota_star)-1);
rkstar = Astar*aalpha*iiota_star*krw_star^(aalpha*iiota_star-1)*kus_star^(aalpha*(1-iiota_star));
rk = A*aalpha*iiota*krw^(aalpha*(1-iiota))*kus^(aalpha*iiota-1);
rkstar = A*aalpha*(1-iiota)*krw^(aalpha*(1-iiota)-1)*kus^(aalpha*iiota);
wstar = Astar*(1-aalpha)*krw_star^(aalpha*iiota_star)*kus_star^(aalpha*(1-iiota_star));
w = A*(1-aalpha)*krw^(aalpha*(1-iiota))*kus^(aalpha*iiota);

0 = -(drdb+nnu*(eeta-1)*brw^(eeta-2))*(c_rw(+1)/c_rw)^(-ggamma)+(1+rb+nnu*brw^(eeta-1))*dMrwdb;
0 = dMrwdb*(rkstar+1-ddelta_rw);

k = kus + kus_star;
kstar = krw + krw_star;

log(nnu) = (1-rrho_nnu)*log(nnu_bar) + rrho_nnu*log(nnu(-1)) + ssigma_nnu*eps_nnu;
log(oomega) = (1-rrho_oomega)*log(oomega_bar) + rrho_oomega*log(oomega(-1)) + ssigma_oomega*eps_oomega;
log(A) = (1-rrho_A)*log(Abar) + rrho_A*log(A(-1)) + ssigma_A*eps_A;
log(Astar) = (1-rrho_Astar)*log(Astarbar) + rrho_Astar*log(Astar(-1)) + ssigma_Astar*eps_A;

vrw = c_rw^(1-ggamma)/(1-ggamma) + bbeta*vrw(+1);
vus = c_us^(1-ggamma)/(1-ggamma) + bbeta*vus(+1);

y = A*(kus^iiota*krw^(1-iiota))^aalpha;

end;

%----------------------------------------------------------------
% 4. Computation
%----------------------------------------------------------------

initval;
  bus = bus_me2;
  brw = brw_me2;
  rb = rb_me2;
  rkstar = rkstar_me2;
  rk = rk_me2;
  krw_star = krw_star_me2;
  kus_star = kus_star_me2;
  krw = krw_me2;
  kus = kus_me2;
  kstar = kstar_me2;
  k = k_me2;
  wstar = wstar_me2;
  w = w_me2;
  c_rw = crw_me2;
  c_us = cus_me2;
  drdb = drdb_me2;
  dMrwdb = dMrwdb_me2;
  nnu = nnu_me2;
  oomega = oomega_me2;
  A = A_me2;
  Astar = Astar_me2;
  vrw = vrw_me2;
  vus = vus_me2;
  y = y_me2;
end;

resid;
check;

shocks;
  var eps_nnu = 1;
  var eps_oomega = 1;
  var eps_A = 1;
end;

set_dynare_seed('default');
stoch_simul(order=2,periods=100000,drop=0,nograph,noprint,pruning);

%----------------------------------------------------------------
% 5. Simulate transition
%----------------------------------------------------------------

spread_path = (rk-ddelta_us-rb)*100;
var_sp = var(spread_path);
auto_sp = autocorr(spread_path);
var_by = var(brw./y);
auto_by = autocorr(brw./y);
corr_pq_by = corr(spread_path,brw./y);
profits = mean((rk-ddelta_us-rb).*(bus) - oomega./(1+llambda).*(bus).^(1+llambda));

moments = [mean(brw./y) mean(spread_path) var_by var_sp corr_pq_by auto_by(2) auto_sp(2) profits]';
data_mom = [0.41 0.62 0.03 0.086 -0.56 0.96 0.70 000]';
rowNames = {'Mean b/y','Mean sp','Var b/y','Var sp','Corr (b/y,sp)','Autocorr b/y','Autocorr sp','Profits'};
colNames = {'Model Moments','Data Moments'};
TableA0 = array2table([moments data_mom],'RowNames',rowNames,'VariableNames',colNames)

%----------------------------------------------------------------
% 6. Calculate welfare from transition
%---------------------------------------------------------------

oo_me2 = oo_;
M_me2 = M_;
options_me2 = options_;  

b_sme2 = mean(brw./y);
bus_sme2 = mean(bus./y);
spread_sme2 = mean(spread_path);
rb_sme2 = mean(rb)*100;

save me2_save oo_me2 vus_me2 vrw_me2 cus_me2 crw_me2 M_me2 options_me2 b_sme2 bus_sme2 spread_sme2 rb_sme2;

load me1_save;

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% Simulated 2nd Order Welfare
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

% Stochastic - lag issue fixed
NN = 100000;
BB = 5000;
NN_BB = NN+BB;

vrw_pos = strmatch('vrw',M_.endo_names,'exact');
vus_pos = strmatch('vus',M_.endo_names,'exact');

% Simulate ME1
shock_matrix = randn(M_.exo_nbr,NN_BB);
sim_eq_me1_ergo = simult_(M_me1,options_me1,oo_me1.dr.ys,oo_me1.dr,shock_matrix',options_me1.order);
sim_eq_me1_ergo = sim_eq_me1_ergo(:,BB+2:end); % burn

% ME1 to N=2
shock_matrix = zeros(M_.exo_nbr,2);
for ii=1:NN
  sim_eq_me1 = simult_(M_me1,options_me1,sim_eq_me1_ergo(:,ii),oo_me1.dr,shock_matrix',options_me1.order);
  vrw_sme(ii) = sim_eq_me1(vrw_pos,2);
  vus_sme(ii) = sim_eq_me1(vus_pos,2);
end
for ii=1:NN
  sim_eq_me2 = simult_(M_me2,options_me2,sim_eq_me1_ergo(:,ii),oo_me2.dr,shock_matrix',options_me2.order);
  vrw_me_me2_sim(ii) = sim_eq_me2(vrw_pos,2);
  vus_me_me2_sim(ii) = sim_eq_me2(vus_pos,2);
end

upsilon_us_transition = mean(((vus_me_me2_sim./vus_sme).^(1/(1-ggamma)) - 1)*100)
upsilon_rw_transition = mean(((vrw_me_me2_sim./vrw_sme).^(1/(1-ggamma)) - 1)*100)

Gamma_me2_us = upsilon_us_transition;
Gamma_me2_rw = upsilon_rw_transition;

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

save me2_save oo_me2 vus_me2 vrw_me2 cus_me2 crw_me2 M_me2 options_me2 b_sme2 bus_sme2 spread_sme2 rb_sme2 Gamma_me2_us Gamma_me2_rw;
